# Copyright (c) 2025 Huawei Technologies Co. Ltd.
# deepinsight is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#          http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
import contextlib
import os
from typing import Optional

from dotenv import load_dotenv
from sqlalchemy import create_engine, text, inspect
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

load_dotenv()


class DatabaseConfig:
    """数据库配置类"""

    def __init__(self, db_type: str = None, connection_string: str = None):
        self.db_type = db_type or os.getenv("DB_TYPE", "sqlite")
        self.connection_string = self._get_connection_string(connection_string)

    def _get_connection_string(self, custom_conn_str: Optional[str]) -> str:
        if custom_conn_str:
            return custom_conn_str

        if self.db_type == "postgresql":
            db_user = os.getenv('POSTGRES_USER', 'postgres')
            db_password = os.getenv('POSTGRES_PASSWORD', '')
            db_host = os.getenv('POSTGRES_HOST', 'localhost')
            db_port = os.getenv('POSTGRES_PORT', '5432')
            db_name = os.getenv('POSTGRES_DB', 'postgres')
            return f"postgresql+psycopg2://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}"

        elif self.db_type == "sqlite":
            return os.getenv("SQLITE_CONN_STR", "sqlite:///./app.db")

        raise ValueError(f"不支持的数据库类型: {self.db_type}")

    def _needs_uuid_extension(self) -> bool:
        return self.db_type == "postgresql"

    def _table_exists(self, table_name: str, engine) -> bool:
        """检查表是否存在"""
        inspector = inspect(engine)
        return inspector.has_table(table_name)

    def _init_postgresql_extensions(self, conn):
        """初始化PostgreSQL扩展"""
        if self._needs_uuid_extension():
            conn.execute(text('CREATE EXTENSION IF NOT EXISTS "uuid-ossp"'))

    def _init_tables(self, conn, engine):
        if self.db_type == "sqlite":
            required_tables = {
                "conversation": """
                  CREATE TABLE IF NOT EXISTS conversation (
                      conversation_id TEXT PRIMARY KEY NOT NULL DEFAULT (lower(hex(randomblob(4))) || '-' || lower(hex(randomblob(2))) || '-4' || substr(lower(hex(randomblob(2))),2) || '-' || substr('89ab', abs(random()) % 4 + 1, 1) || substr(lower(hex(randomblob(2))),2) || '-' || lower(hex(randomblob(6)))),
                      user_id TEXT NOT NULL,
                      created_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
                      title TEXT DEFAULT '新建对话',
                      status TEXT NOT NULL DEFAULT 'active'
                  );
                  """,
                "message": """
                  CREATE TABLE IF NOT EXISTS message (
                      message_id TEXT PRIMARY KEY NOT NULL DEFAULT (lower(hex(randomblob(4))) || '-' || lower(hex(randomblob(2))) || '-4' || substr(lower(hex(randomblob(2))),2) || '-' || substr('89ab', abs(random()) % 4 + 1, 1) || substr(lower(hex(randomblob(2))),2) || '-' || lower(hex(randomblob(6)))),
                      conversation_id TEXT NOT NULL,
                      content TEXT NOT NULL,
                      type TEXT NOT NULL,
                      created_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
                      FOREIGN KEY (conversation_id) REFERENCES conversation(conversation_id) ON DELETE CASCADE
                  );
                  """,
                "report": """
                  CREATE TABLE IF NOT EXISTS report (
                      report_id TEXT PRIMARY KEY NOT NULL DEFAULT (lower(hex(randomblob(4))) || '-' || lower(hex(randomblob(2))) || '-4' || substr(lower(hex(randomblob(2))),2) || '-' || substr('89ab', abs(random()) % 4 + 1, 1) || substr(lower(hex(randomblob(2))),2) || '-' || lower(hex(randomblob(6)))),
                      message_id TEXT NOT NULL,
                      conversation_id TEXT NOT NULL,
                      thought TEXT,
                      report_content TEXT NOT NULL,
                      created_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
                      FOREIGN KEY (message_id) REFERENCES message(message_id) ON DELETE CASCADE,
                      FOREIGN KEY (conversation_id) REFERENCES conversation(conversation_id) ON DELETE CASCADE
                  );
                  """
            }
        else:
            required_tables = {
                "conversation": """
                CREATE TABLE IF NOT EXISTS conversation (
                    conversation_id UUID PRIMARY KEY NOT NULL DEFAULT uuid_generate_v4(),
                    user_id VARCHAR(36) NOT NULL,
                    created_time TIMESTAMP(3) WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
                    title VARCHAR(255) DEFAULT '新建对话',
                    status VARCHAR(50) NOT NULL DEFAULT 'active'
                );
                """,
                "message": """
                CREATE TABLE IF NOT EXISTS message (
                    message_id UUID PRIMARY KEY NOT NULL DEFAULT uuid_generate_v4(),
                    conversation_id UUID NOT NULL,
                    content TEXT NOT NULL,
                    type VARCHAR(50) NOT NULL,
                    created_time TIMESTAMP(3) WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
                    FOREIGN KEY (conversation_id) REFERENCES conversation(conversation_id) ON DELETE CASCADE
                );
                """,
                "report": """
                CREATE TABLE IF NOT EXISTS report (
                    report_id UUID PRIMARY KEY NOT NULL DEFAULT uuid_generate_v4(),
                    message_id UUID NOT NULL,
                    conversation_id UUID NOT NULL,
                    thought TEXT,
                    report_content TEXT NOT NULL,
                    created_time TIMESTAMP(3) WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
                    FOREIGN KEY (message_id) REFERENCES message(message_id) ON DELETE CASCADE,
                    FOREIGN KEY (conversation_id) REFERENCES conversation(conversation_id) ON DELETE CASCADE
                );
                """
            }

        for table_name, ddl in required_tables.items():
            if not self._table_exists(table_name, engine):
                print(f"创建表 {table_name}...")
                conn.execute(text(ddl))

    def initialize(self, engine):
        """执行数据库初始化"""
        try:
            with engine.begin() as conn:
                if self._needs_uuid_extension():
                    self._init_postgresql_extensions(conn)
                self._init_tables(conn, engine)
            print("✅ 数据库初始化完成")
            return True
        except SQLAlchemyError as e:
            print(f"❌ 数据库初始化失败: {str(e)}")
            raise

# 创建默认数据库配置
default_config = DatabaseConfig()

# 创建引擎
engine = create_engine(
    default_config.connection_string,
    echo=False,  # 设置为True可打印SQL语句，调试时使用
    connect_args={"check_same_thread": False} if default_config.db_type == "sqlite" else {}
)

default_config.initialize(engine)

# SQLAlchemy模型基类
# 所有数据模型都应继承此类，用于表结构定义
DatabaseModel = declarative_base()

# 数据库会话工厂
DatabaseSession = sessionmaker(autocommit=False, autoflush=False, bind=engine)


@contextlib.contextmanager
def get_database_session():
    """
    获取数据库会话的生成器函数

    Yields:
        Session: 数据库会话对象

    Note:
        使用FastAPI等框架时，可作为依赖项注入
        会话会在使用后自动关闭
    """
    session = DatabaseSession()
    try:
        yield session
    finally:
        session.close()
